#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
code_mapping_analyzer.py

编码转换辅助分析工具：分析两个 TIF 文件中编码的对应关系，生成 Excel 报告。

功能：
- 读取两个 TIF 文件（a.tif 和 b.tif）
- 分析 a.tif 中每个代码在 b.tif 中的分布情况
- 生成详细的 Excel 报告，包含：
  * 代码对应统计表
  * 每个 a.tif 代码在 b.tif 中的频次分布
  * 建议的转换映射（基于最大频次）
  * 重叠区域分析
- 支持指定分析区域（通过边界框或掩膜文件）
- 支持忽略 nodata 值

依赖：rasterio, numpy, pandas, openpyxl

用法示例：
1) 基本分析：
   python code_mapping_analyzer.py --a-tif a.tif --b-tif b.tif --output mapping_analysis.xlsx

2) 指定分析区域：
   python code_mapping_analyzer.py --a-tif a.tif --b-tif b.tif --output mapping_analysis.xlsx --bounds 120 30 130 40

3) 使用掩膜文件：
   python code_mapping_analyzer.py --a-tif a.tif --b-tif b.tif --output mapping_analysis.xlsx --mask mask.tif
"""

import argparse
import os
from typing import Dict, List, Optional, Tuple, Union

import numpy as np
import pandas as pd
import rasterio
from rasterio.warp import transform_bounds


class CodeMappingAnalyzer:
    """编码映射分析器"""
    
    def __init__(self, a_tif: str, b_tif: str, bounds: Optional[Tuple[float, float, float, float]] = None, 
                 mask_tif: Optional[str] = None, ignore_nodata: bool = True):
        """
        初始化分析器
        
        Args:
            a_tif: 源 TIF 文件路径
            b_tif: 目标 TIF 文件路径
            bounds: 分析区域边界框 (minx, miny, maxx, maxy)
            mask_tif: 掩膜 TIF 文件路径（1=分析区域，0=忽略）
            ignore_nodata: 是否忽略 nodata 值
        """
        self.a_tif = a_tif
        self.b_tif = b_tif
        self.bounds = bounds
        self.mask_tif = mask_tif
        self.ignore_nodata = ignore_nodata
        
        # 打开文件
        self.a_src = rasterio.open(a_tif)
        self.b_src = rasterio.open(b_tif)
        
        # 确保两个文件在同一坐标系
        if self.a_src.crs != self.b_src.crs:
            print(f"警告: 两个文件的 CRS 不同")
            print(f"  a.tif CRS: {self.a_src.crs}")
            print(f"  b.tif CRS: {self.b_src.crs}")
            print("将使用 a.tif 的 CRS 作为参考")
    
    def get_analysis_window(self) -> Tuple[rasterio.windows.Window, rasterio.windows.Window]:
        """获取分析窗口"""
        if self.bounds:
            # 使用指定的边界框
            minx, miny, maxx, maxy = self.bounds
            a_window = rasterio.windows.from_bounds(minx, miny, maxx, maxy, self.a_src.transform)
            b_window = rasterio.windows.from_bounds(minx, miny, maxx, maxy, self.b_src.transform)
        else:
            # 使用两个文件的交集
            a_bounds = self.a_src.bounds
            b_bounds = self.b_src.bounds
            
            # 计算交集
            minx = max(a_bounds.left, b_bounds.left)
            miny = max(a_bounds.bottom, b_bounds.bottom)
            maxx = min(a_bounds.right, b_bounds.right)
            maxy = min(a_bounds.top, b_bounds.top)
            
            if minx >= maxx or miny >= maxy:
                raise ValueError("两个文件没有重叠区域")
            
            a_window = rasterio.windows.from_bounds(minx, miny, maxx, maxy, self.a_src.transform)
            b_window = rasterio.windows.from_bounds(minx, miny, maxx, maxy, self.b_src.transform)
        
        return a_window, b_window
    
    def read_data(self) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
        """读取分析数据"""
        a_window, b_window = self.get_analysis_window()
        
        # 读取数据
        a_data = self.a_src.read(1, window=a_window)  # 只读第一个波段
        b_data = self.b_src.read(1, window=b_window)
        
        # 如果两个数据形状不同，需要重采样到相同网格
        if a_data.shape != b_data.shape:
            print(f"数据形状不同，正在重采样对齐...")
            print(f"  a.tif 形状: {a_data.shape}")
            print(f"  b.tif 形状: {b_data.shape}")
            
            # 使用 a.tif 的网格作为参考
            target_transform = self.a_src.window_transform(a_window)
            target_crs = self.a_src.crs
            
            # 重采样 b.tif 到 a.tif 的网格
            # 使用 b.tif 的原始数据类型，避免数据截断
            b_dtype = self.b_src.dtypes[0]  # 获取 b.tif 的原始数据类型
            b_resampled = np.empty(a_data.shape, dtype=b_dtype)
            from rasterio.warp import reproject
            reproject(
                b_data, b_resampled,
                src_transform=self.b_src.window_transform(b_window),
                dst_transform=target_transform,
                src_crs=self.b_src.crs,
                dst_crs=target_crs,
                resampling=rasterio.enums.Resampling.nearest
            )
            b_data = b_resampled
        
        # 读取掩膜（如果提供）
        mask_data = None
        if self.mask_tif:
            with rasterio.open(self.mask_tif) as mask_src:
                mask_window = rasterio.windows.from_bounds(
                    *rasterio.windows.bounds(a_window, self.a_src.transform),
                    mask_src.transform
                )
                mask_data = mask_src.read(1, window=mask_window)
                # 确保掩膜与数据尺寸一致
                if mask_data.shape != a_data.shape:
                    from rasterio.warp import reproject
                    mask_resampled = np.empty_like(a_data)
                    reproject(
                        mask_data, mask_resampled,
                        src_transform=mask_src.transform,
                        dst_transform=self.a_src.window_transform(a_window),
                        src_crs=mask_src.crs,
                        dst_crs=self.a_src.crs,
                        resampling=rasterio.enums.Resampling.nearest
                    )
                    mask_data = mask_resampled
        
        return a_data, b_data, mask_data
    
    def analyze_mapping(self) -> Dict:
        """分析编码映射关系"""
        a_data, b_data, mask_data = self.read_data()
        
        print(f"原始数据形状: a_data={a_data.shape}, b_data={b_data.shape}")
        print(f"a.tif 数据类型: {a_data.dtype}, 数据范围: {np.min(a_data)} - {np.max(a_data)}")
        print(f"b.tif 数据类型: {b_data.dtype}, 数据范围: {np.min(b_data)} - {np.max(b_data)}")
        print(f"a.tif nodata: {self.a_src.nodata}")
        print(f"b.tif nodata: {self.b_src.nodata}")
        
        # 应用掩膜
        if mask_data is not None:
            valid_mask = mask_data == 1
            print(f"掩膜应用前有效像素: {np.sum(valid_mask)}")
            a_data = a_data[valid_mask]
            b_data = b_data[valid_mask]
        else:
            valid_mask = np.ones_like(a_data, dtype=bool)
        
        # 智能处理 nodata
        if self.ignore_nodata:
            # 检查 a.tif 的 nodata 值是否合理
            a_nodata = self.a_src.nodata
            if a_nodata is not None:
                a_nodata_count = np.sum(a_data == a_nodata)
                a_total_count = a_data.size
                a_nodata_ratio = a_nodata_count / a_total_count
                print(f"a.tif nodata={a_nodata}, 占比={a_nodata_ratio:.2%}")
                
                # 如果 nodata 占比过高（>90%），可能不是真正的 nodata
                if a_nodata_ratio > 0.9:
                    print(f"警告: a.tif 的 nodata 值占比过高，可能不是真正的 nodata，将忽略此设置")
                    a_nodata = None
                else:
                    valid_mask &= (a_data != a_nodata)
                    print(f"忽略 a.tif nodata 后有效像素: {np.sum(valid_mask)}")
            
            # 检查 b.tif 的 nodata 值是否合理
            b_nodata = self.b_src.nodata
            if b_nodata is not None:
                b_nodata_count = np.sum(b_data == b_nodata)
                b_total_count = b_data.size
                b_nodata_ratio = b_nodata_count / b_total_count
                print(f"b.tif nodata={b_nodata}, 占比={b_nodata_ratio:.2%}")
                
                # 如果 nodata 占比过高（>90%），可能不是真正的 nodata
                if b_nodata_ratio > 0.9:
                    print(f"警告: b.tif 的 nodata 值占比过高，可能不是真正的 nodata，将忽略此设置")
                    b_nodata = None
                else:
                    valid_mask &= (b_data != b_nodata)
                    print(f"忽略 b.tif nodata 后有效像素: {np.sum(valid_mask)}")
            
            a_data = a_data[valid_mask]
            b_data = b_data[valid_mask]
        
        print(f"最终有效像素数: {len(a_data)}")
        print(f"a.tif 唯一值数量: {len(np.unique(a_data))}")
        print(f"b.tif 唯一值数量: {len(np.unique(b_data))}")
        
        # 统计对应关系
        mapping_stats = {}
        a_codes = np.unique(a_data)
        
        for a_code in a_codes:
            a_mask = (a_data == a_code)
            b_values = b_data[a_mask]
            
            if len(b_values) == 0:
                continue
            
            # 统计 b.tif 中的值分布
            unique_b, counts = np.unique(b_values, return_counts=True)
            total_pixels = len(b_values)
            
            # 计算百分比
            percentages = (counts / total_pixels) * 100
            
            # 找到最频繁的 b 值
            max_idx = np.argmax(counts)
            most_frequent_b = unique_b[max_idx]
            most_frequent_count = counts[max_idx]
            most_frequent_pct = percentages[max_idx]
            
            mapping_stats[a_code] = {
                'total_pixels': total_pixels,
                'b_values': unique_b.tolist(),
                'b_counts': counts.tolist(),
                'b_percentages': percentages.tolist(),
                'most_frequent_b': most_frequent_b,
                'most_frequent_count': most_frequent_count,
                'most_frequent_pct': most_frequent_pct,
                'b_value_distribution': dict(zip(unique_b.tolist(), percentages.tolist()))
            }
        
        return mapping_stats
    
    def generate_excel_report(self, output_path: str, mapping_stats: Dict):
        """生成 Excel 报告"""
        with pd.ExcelWriter(output_path, engine='openpyxl') as writer:
            # 检查是否有数据
            if not mapping_stats:
                # 创建空报告
                empty_df = pd.DataFrame({'说明': ['未找到有效的编码对应关系', '可能原因：', '1. 两个文件没有重叠区域', '2. 所有数据都是 nodata 值', '3. 文件格式或读取问题']})
                empty_df.to_excel(writer, sheet_name='错误信息', index=False)
                return
            
            # 1. 代码对应统计表
            summary_data = []
            for a_code, stats in mapping_stats.items():
                summary_data.append({
                    'A_TIF代码': a_code,
                    '总像素数': stats['total_pixels'],
                    'B_TIF中最频繁代码': stats['most_frequent_b'],
                    '最频繁代码像素数': stats['most_frequent_count'],
                    '最频繁代码占比(%)': round(stats['most_frequent_pct'], 2),
                    'B_TIF中不同代码数': len(stats['b_values'])
                })
            
            summary_df = pd.DataFrame(summary_data)
            summary_df.to_excel(writer, sheet_name='代码对应统计', index=False)
            
            # 2. 详细分布表
            detail_data = []
            for a_code, stats in mapping_stats.items():
                for b_code, count, pct in zip(stats['b_values'], stats['b_counts'], stats['b_percentages']):
                    detail_data.append({
                        'A_TIF代码': a_code,
                        'B_TIF代码': b_code,
                        '像素数': count,
                        '占比(%)': round(pct, 2)
                    })
            
            if detail_data:
                detail_df = pd.DataFrame(detail_data)
                detail_df = detail_df.sort_values(['A_TIF代码', '像素数'], ascending=[True, False])
                detail_df.to_excel(writer, sheet_name='详细分布', index=False)
            else:
                empty_detail_df = pd.DataFrame({'说明': ['无详细分布数据']})
                empty_detail_df.to_excel(writer, sheet_name='详细分布', index=False)
            
            # 3. 建议转换映射
            mapping_data = []
            for a_code, stats in mapping_stats.items():
                mapping_data.append({
                    'A_TIF代码': a_code,
                    '建议B_TIF代码': stats['most_frequent_b'],
                    '置信度(%)': round(stats['most_frequent_pct'], 2),
                    '总像素数': stats['total_pixels']
                })
            
            mapping_df = pd.DataFrame(mapping_data)
            mapping_df.to_excel(writer, sheet_name='建议转换映射', index=False)
            
            # 4. 数据质量评估
            quality_data = []
            total_a_pixels = sum(stats['total_pixels'] for stats in mapping_stats.values())
            
            for a_code, stats in mapping_stats.items():
                pct_of_total = (stats['total_pixels'] / total_a_pixels) * 100
                quality_data.append({
                    'A_TIF代码': a_code,
                    '像素数': stats['total_pixels'],
                    '占总像素比例(%)': round(pct_of_total, 2),
                    'B_TIF代码多样性': len(stats['b_values']),
                    '最大占比(%)': round(stats['most_frequent_pct'], 2),
                    '映射质量': '高' if stats['most_frequent_pct'] > 80 else '中' if stats['most_frequent_pct'] > 50 else '低'
                })
            
            quality_df = pd.DataFrame(quality_data)
            quality_df.to_excel(writer, sheet_name='数据质量评估', index=False)
    
    def close(self):
        """关闭文件"""
        self.a_src.close()
        self.b_src.close()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description='分析两个 TIF 文件的编码对应关系')
    
    # 必需参数
    parser.add_argument('--a-tif', required=True, help='源 TIF 文件路径')
    parser.add_argument('--b-tif', required=True, help='目标 TIF 文件路径')
    parser.add_argument('--output', required=True, help='输出 Excel 文件路径')
    
    # 可选参数
    parser.add_argument('--bounds', nargs=4, type=float, metavar=('MINX', 'MINY', 'MAXX', 'MAXY'),
                       help='分析区域边界框 (minx, miny, maxx, maxy)')
    parser.add_argument('--mask', help='掩膜 TIF 文件路径（1=分析区域，0=忽略）')
    parser.add_argument('--no-ignore-nodata', action='store_true',
                       help='不忽略 nodata 值（默认忽略）')
    
    return parser.parse_args()


def main():
    args = parse_args()
    
    # 验证输入文件
    if not os.path.exists(args.a_tif):
        raise SystemExit(f"文件不存在: {args.a_tif}")
    if not os.path.exists(args.b_tif):
        raise SystemExit(f"文件不存在: {args.b_tif}")
    if args.mask and not os.path.exists(args.mask):
        raise SystemExit(f"掩膜文件不存在: {args.mask}")
    
    # 创建输出目录
    os.makedirs(os.path.dirname(args.output), exist_ok=True)
    
    # 解析边界框
    bounds = None
    if args.bounds:
        bounds = tuple(args.bounds)
    
    # 创建分析器
    analyzer = CodeMappingAnalyzer(
        a_tif=args.a_tif,
        b_tif=args.b_tif,
        bounds=bounds,
        mask_tif=args.mask,
        ignore_nodata=not args.no_ignore_nodata
    )
    
    try:
        print("正在分析编码对应关系...")
        mapping_stats = analyzer.analyze_mapping()
        
        print(f"找到 {len(mapping_stats)} 个 A_TIF 代码")
        print("正在生成 Excel 报告...")
        analyzer.generate_excel_report(args.output, mapping_stats)
        
        print(f"分析完成！报告已保存到: {args.output}")
        
        # 显示简要统计
        print("\n简要统计:")
        for a_code, stats in list(mapping_stats.items())[:5]:  # 只显示前5个
            print(f"  A代码 {a_code} -> B代码 {stats['most_frequent_b']} "
                  f"(置信度: {stats['most_frequent_pct']:.1f}%)")
        if len(mapping_stats) > 5:
            print(f"  ... 还有 {len(mapping_stats) - 5} 个代码")
            
    except Exception as e:
        print(f"分析过程中发生错误: {e}")
        raise
    finally:
        analyzer.close()


if __name__ == '__main__':
    main()
